# -*- coding: utf-8 -*-
"""
数据获取模块 - 从多个数据源获取股票数据
支持akshare、yfinance等数据源
"""
import akshare as ak
import yfinance as yf
import pandas as pd
import numpy as np
import time
import requests
from datetime import datetime, timedelta
import logging
from typing import Optional, Dict, List, Tuple
import warnings
warnings.filterwarnings('ignore')

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class StockDataFetcher:
    """股票数据获取器"""
    
    def __init__(self):
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
    
    def get_stock_list(self) -> pd.DataFrame:
        """获取A股股票列表"""
        try:
            logger.info("正在获取A股股票列表...")
            # 获取沪深A股股票列表
            stock_list_sh = ak.stock_info_a_code_name()
            time.sleep(1)
            
            # 添加市场标识
            stock_list_sh['market'] = stock_list_sh['code'].apply(
                lambda x: 'SH' if x.startswith(('60', '68', '11', '13')) else 'SZ'
            )
            
            logger.info(f"成功获取 {len(stock_list_sh)} 只股票信息")
            return stock_list_sh
            
        except Exception as e:
            logger.error(f"获取股票列表失败: {e}")
            return pd.DataFrame()
    
    def get_stock_basic_info(self, stock_code: str) -> Dict:
        """获取股票基本信息"""
        try:
            # 获取股票基本信息
            basic_info = ak.stock_individual_info_em(symbol=stock_code)
            
            # 获取财务指标
            financial_info = ak.stock_financial_em(symbol=stock_code)
            
            # 获取分红信息
            dividend_info = self.get_dividend_info(stock_code)
            
            result = {
                'basic': basic_info,
                'financial': financial_info,
                'dividend': dividend_info
            }
            
            return result
            
        except Exception as e:
            logger.warning(f"获取股票 {stock_code} 基本信息失败: {e}")
            return {}
    
    def get_dividend_info(self, stock_code: str) -> pd.DataFrame:
        """获取分红信息"""
        try:
            # 使用正确的akshare函数获取分红信息
            dividend_df = ak.stock_history_dividend_detail(symbol=stock_code)
            return dividend_df
        except Exception as e:
            logger.warning(f"获取股票 {stock_code} 分红信息失败: {e}")
            return pd.DataFrame()
    
    def get_stock_realtime_data(self, stock_codes: List[str]) -> pd.DataFrame:
        """获取股票实时数据"""
        try:
            # 使用akshare获取实时数据
            realtime_data = ak.stock_zh_a_spot_em()
            
            # 筛选指定股票
            if stock_codes:
                realtime_data = realtime_data[realtime_data['代码'].isin(stock_codes)]
            
            return realtime_data
            
        except Exception as e:
            logger.error(f"获取实时数据失败: {e}")
            return pd.DataFrame()
    
    def get_stock_historical_data(self, stock_code: str, start_date: str = None, 
                                end_date: str = None, period: str = "daily") -> pd.DataFrame:
        """获取股票历史数据"""
        try:
            if start_date is None:
                start_date = (datetime.now() - timedelta(days=365*2)).strftime("%Y%m%d")
            if end_date is None:
                end_date = datetime.now().strftime("%Y%m%d")
            
            # 使用akshare获取历史数据
            hist_data = ak.stock_zh_a_hist(
                symbol=stock_code, 
                period=period,
                start_date=start_date,
                end_date=end_date,
                adjust="qfq"  # 前复权
            )
            
            if not hist_data.empty:
                # 先检查数据的实际列数和列名
                logger.info(f"原始数据列数: {len(hist_data.columns)}, 列名: {list(hist_data.columns)}")
                
                # 动态处理列名重命名，防止长度不匹配
                expected_columns = ['date', 'stock_code', 'open', 'close', 'high', 'low', 'volume', 'turnover', 'amplitude', 'change_pct', 'change_amount', 'turnover_rate']
                
                if len(hist_data.columns) == len(expected_columns):
                    # 如果列数匹配，直接重命名
                    hist_data.columns = expected_columns
                elif len(hist_data.columns) == len(expected_columns) - 1:
                    # 如果少一列（可能没有stock_code列），去掉stock_code
                    hist_data.columns = ['date', 'open', 'close', 'high', 'low', 'volume', 'turnover', 'amplitude', 'change_pct', 'change_amount', 'turnover_rate']
                else:
                    # 使用原始列名，但至少确保有基础列
                    logger.warning(f"列数不匹配，使用原始列名: {list(hist_data.columns)}")
                    # 创建列名映射
                    column_mapping = {}
                    if '日期' in hist_data.columns:
                        column_mapping['日期'] = 'date'
                    if '开盘' in hist_data.columns:
                        column_mapping['开盘'] = 'open'
                    if '收盘' in hist_data.columns:
                        column_mapping['收盘'] = 'close'
                    if '最高' in hist_data.columns:
                        column_mapping['最高'] = 'high'
                    if '最低' in hist_data.columns:
                        column_mapping['最低'] = 'low'
                    if '成交量' in hist_data.columns:
                        column_mapping['成交量'] = 'volume'
                    
                    hist_data = hist_data.rename(columns=column_mapping)
                
                # 确保有date列用于设置索引
                if 'date' in hist_data.columns:
                    hist_data['date'] = pd.to_datetime(hist_data['date'])
                    hist_data.set_index('date', inplace=True)
                elif '日期' in hist_data.columns:
                    hist_data['日期'] = pd.to_datetime(hist_data['日期'])
                    hist_data.set_index('日期', inplace=True)
                
                hist_data = hist_data.sort_index()
                
            return hist_data
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 历史数据失败: {e}")
            return pd.DataFrame()
    
    def get_market_data(self) -> Dict:
        """获取大盘数据"""
        try:
            # 获取上证指数
            sh_index = ak.stock_zh_index_spot_em(symbol="000001")
            
            # 获取深证成指
            sz_index = ak.stock_zh_index_spot_em(symbol="399001")
            
            # 获取创业板指
            cy_index = ak.stock_zh_index_spot_em(symbol="399006")
            
            market_data = {
                'shanghai': sh_index,
                'shenzhen': sz_index,
                'chinext': cy_index
            }
            
            return market_data
            
        except Exception as e:
            logger.error(f"获取大盘数据失败: {e}")
            return {}
    
    def get_industry_data(self, stock_code: str) -> Dict:
        """获取行业板块数据"""
        try:
            # 获取股票所属行业
            stock_industry = ak.stock_board_industry_name_em()
            
            # 获取概念板块
            stock_concept = ak.stock_board_concept_name_em()
            
            industry_data = {
                'industry': stock_industry,
                'concept': stock_concept
            }
            
            return industry_data
            
        except Exception as e:
            logger.warning(f"获取行业数据失败: {e}")
            return {}
    
    def validate_stock_code(self, stock_code: str) -> bool:
        """验证股票代码是否有效"""
        try:
            data = self.get_stock_historical_data(stock_code, 
                start_date=(datetime.now() - timedelta(days=5)).strftime("%Y%m%d"))
            return not data.empty
        except:
            return False
    
    def get_stock_financial_indicators(self, stock_code: str) -> Dict:
        """获取股票财务指标"""
        try:
            # 获取主要财务指标
            financial_data = ak.stock_financial_abstract_ths(symbol=stock_code)
            
            # 获取现金流量表
            cashflow_data = ak.stock_cashflow_ths(symbol=stock_code)
            
            # 获取资产负债表
            balance_data = ak.stock_balance_sheet_by_yearly_ths(symbol=stock_code)
            
            result = {
                'financial': financial_data,
                'cashflow': cashflow_data,
                'balance': balance_data
            }
            
            return result
            
        except Exception as e:
            logger.warning(f"获取股票 {stock_code} 财务指标失败: {e}")
            return {}
    
    def get_stock_data(self, stock_code: str, period: str = "1y") -> Optional[pd.DataFrame]:
        """获取股票数据的统一接口"""
        try:
            # 转换period参数
            if period == "1y":
                start_date = (datetime.now() - timedelta(days=365)).strftime("%Y%m%d")
            elif period == "6mo":
                start_date = (datetime.now() - timedelta(days=180)).strftime("%Y%m%d")
            elif period == "3mo":
                start_date = (datetime.now() - timedelta(days=90)).strftime("%Y%m%d")
            elif period == "1mo":
                start_date = (datetime.now() - timedelta(days=30)).strftime("%Y%m%d")
            else:
                start_date = (datetime.now() - timedelta(days=365)).strftime("%Y%m%d")
            
            end_date = datetime.now().strftime("%Y%m%d")
            
            return self.get_stock_historical_data(stock_code, start_date, end_date)
            
        except Exception as e:
            logger.error(f"获取股票数据失败: {e}")
            return None

# 单例模式
_data_fetcher = None

def get_data_fetcher() -> StockDataFetcher:
    """获取数据获取器实例"""
    global _data_fetcher
    if _data_fetcher is None:
        _data_fetcher = StockDataFetcher()
    return _data_fetcher

if __name__ == "__main__":
    # 测试数据获取
    fetcher = get_data_fetcher()
    
    # 测试获取股票列表
    stock_list = fetcher.get_stock_list()
    print(f"股票列表数量: {len(stock_list)}")
    
    # 测试获取历史数据
    if not stock_list.empty:
        test_code = stock_list.iloc[0]['code']
        hist_data = fetcher.get_stock_historical_data(test_code)
        print(f"测试股票 {test_code} 历史数据: {len(hist_data)} 条")
